#!/bin/bash

######################################################################################
# 网卡中断绑核脚本：一键查看或操作网卡中断绑核
#
# 使用方式：
#     查看绑核情况：./nic_affinity.sh -i enp5s0
#     进行中断绑核：./nic_affinity.sh -i enp5s0 -c 0-63
#     更多使用方式：./nic_affinity.sh --help
#
# 常见问题：
#     脚本从Windows放到Linux环境后，可能有字符问题，需要修复：dos2unix nic_affinity.sh
######################################################################################

show_affinity() {
    local nic=$1
    local pcie=$(ethtool -i ${nic} | grep bus-info | awk '{print $2}')

    local irq_grep_pattern="${nic}|${pcie}"
    if dmesg | grep renamed | grep ${pcie} &>/dev/null ; then
        local rename="$(dmesg | grep renamed | grep ${pcie} | awk '{print $NF}')"
        irq_grep_pattern="${irq_grep_pattern}|${rename}"
    fi

    # 排除非收发包队列
    local irq_grep_except="mlx5_pages_eq|mlx5_cmd_eq|mlx5_async|hclge-misc|hinic_aeq|hinic_ceq"

    local nic_numa=$(lspci -vvv -s ${pcie} | grep -i numa | awk '{print $NF}')

    # 打印网卡信息
    echo "NIC: ${nic}"
    echo "Bus Info: ${pcie}"
    echo "NUMA: ${nic_numa}"
    echo -e "\nCurrent affinity:"
    for irq in $(cat /proc/interrupts | grep -E "$irq_grep_pattern" | grep -vE "$irq_grep_except" | awk -F ':' '{print $1}'); do
        echo "irq: ${irq} -> core: $(cat /proc/irq/${irq}/smp_affinity_list)"
    done
    echo
}

set_affinity() {
    local nic=$1
    local cpu_list=$2
    local range=$3

    local cpu_array=($(parse_cpu_list $2))

    local pcie=$(ethtool -i ${nic} | grep bus-info | awk '{print $2}')

    local irq_grep_pattern="${nic}|${pcie}"
    if dmesg | grep renamed | grep ${pcie} &>/dev/null ; then
        local rename="$(dmesg | grep renamed | grep ${pcie} | awk '{print $NF}')"
        irq_grep_pattern="${irq_grep_pattern}|${rename}"
    fi

    # 排除非收发包队列
    local irq_grep_except="mlx5_pages_eq|mlx5_cmd_eq|mlx5_async|hclge-misc|hinic_aeq|hinic_ceq"

    if [ $range -eq 0 ]; then
        local num_cpus=${#cpu_array[@]}
        local i=0

        for irq in $(cat /proc/interrupts | grep -E "$irq_grep_pattern" | grep -vE "$irq_grep_except" | awk -F ':' '{print $1}'); do
            echo ${cpu_array[i]} > /proc/irq/${irq}/smp_affinity_list
            i=$(((i + 1) % ${num_cpus})) # 若CPU不够，则循环绑核
        done

        if systemctl status irqbalance | grep Active | grep running &>/dev/null ; then
            echo "Warning: irqbalance service is still running, please consider to stop it."
        fi
    else
        for irq in $(cat /proc/interrupts | grep -E "$irq_grep_pattern" | grep -vE "$irq_grep_except" | awk -F ':' '{print $1}'); do
            echo $cpu_list > /proc/irq/${irq}/smp_affinity_list
        done
    fi
}

parse_cpu_list() {
    local cpu_list=$1

    IFS=, read -ra parts <<< "${cpu_list}" # 分割字符串到数组

    local cpu_array=()
    for part in "${parts[@]}"; do
        if [[ $part == *"-"* ]]; then
            IFS=- read start end <<< "$part"
            if (( start > end )); then
                echo "Error: Invalid range ${start}-${end}" >&2
                return 1
            fi
            for (( i=start; i<=end; i++ )); do
                cpu_array+=("$i") # 生成范围内的所有数字
            done
        else
            cpu_array+=("$part")
        fi
    done

    # 去重并排序
    local sorted_cpus=($(printf "%s\n" "${cpu_array[@]}" | sort -nu | tr '\n' ' '))
    echo "${sorted_cpus[@]}"

    return 0
}

convert_node_list_to_cpu_list() {
    local node_list=$1
    IFS=, read -ra parts <<< "${node_list}" # 分割字符串到数组

    local node_array=()
    for part in "${parts[@]}"; do
        if [[ $part == *"-"* ]]; then
            IFS=- read start end <<< "$part"
            if (( start > end )); then
                echo "Error: Invalid range ${start}-${end}" >&2
                return 1
            fi
            for (( i=start; i<=end; i++ )); do
                node_array+=("$i") # 生成范围内的所有数字
            done
        else
            node_array+=("$part")
        fi
    done

    local core_list=""
    for i in "${node_array[@]}"; do
        local cpu_list_file="/sys/devices/system/node/node${i}/cpulist"

        if [[ ! -f "$cpu_list_file" ]]; then
            echo "Error: NUMA node $i does not exist" >&2
            return 1
        fi

        core_list="$core_list,$(cat $cpu_list_file)"
    done

    core_list="${core_list#,}" # 去除逗号前缀
    echo $core_list

    return 0
}

help() {
    cat <<- EOF

Usage:
    Display current affinity: $0 -i NIC
    Set new affinity:         $0 -i NIC {-c CPU_LIST | -n NODE_LIST} [--range]

Options:
    -i, --interface     Network interface. Example: enp5s0
    -c, --cpu-list      CPU list to bind. Example: 0-19,40-59
    -n, --node-list     NUMA node list to bind. Example: 0,1.
    -r, --range         Bind all queues to the same range specified by the CPU list instead of single core
    -h, --help          Show usage.

Example:
    查看绑核情况
    $0 -i enp5s0

    执行中断绑核，中断和核一一对应
    $0 -i enp5s0 -c 0-63

    执行中断绑核，中断绑定指定范围的核
    $0 -i enp5s0 -c 0-63 --range

    执行中断绑核，中断绑定指定NUMA节点下的核
    $0 -i enp5s0 -n 0

EOF
}

main() {
    local nic=""
    local cpu_list=""
    local node_list=""
    local range=0

    # 检查GNU getopt兼容性
    getopt -T >/dev/null 2>&1
    local ret=$?
    if [ $ret -ne 4 ]; then
        echo "Error: This script requires GNU getopt for long option support." >&2
        exit 1
    fi

    # 定义短选项和长选项
    local short_opts="i:c:n:rh"
    local long_opts="interface:,cpu-list:,node-list:,range,help"

    # 使用getopt解析参数
    local parsed
    parsed=$(getopt -o "$short_opts" -l "$long_opts" -n "$0" -- "$@")

    if [ $? -ne 0 ]; then
        help
        exit 1
    fi
    eval set -- "$parsed"

    # 处理解析后的选项
    while true; do
        case "$1" in
            -i | --interface)
                nic="$2"
                shift 2
                ;;
            -c | --cpu-list)
                cpu_list="$2"
                shift 2
                ;;
            -n | --node-list)
                node_list="$2"
                shift 2
                ;;
            -r | --range)
                range=1
                shift 1
                ;;
            -h|--help)
                help
                exit 0
                ;;
            --)
                shift
                break
                ;;
            *)
                echo "Error: Internal error!" >&2
                exit 1
                ;;
        esac
    done

    # 检查是否已传入网口名
    if [ -z "${nic}" ]; then
        echo "Error: Missing network interface name." >&2
        help
        exit 1
    fi

    # 检查网口是否存在
    local all_nics=" $(cat /proc/net/dev | awk -F ':' '/:/{ORS=" "; print $1}') "
    if ! echo "${all_nics}" | grep " ${nic} " &>/dev/null ; then
        echo "Error: Invalid network interface: ${nic}" >&2
        exit 1
    fi

    # 若没有传入CPU或NUMA列表，则显示当前绑核情况
    if [ -z "${cpu_list}" ] && [ -z "${node_list}" ]; then
        show_affinity ${nic}
        exit 0
    fi

    # 不能同时指定CPU和NUMA列表
    if [ -n "${cpu_list}" ] && [ -n "${node_list}" ]; then
        echo "Error: Cannot specify '-c' and '-n' simultaneously" >&2
        exit 1
    fi

    # 按CPU列表绑核
    if [ -n "${cpu_list}" ]; then
        set_affinity ${nic} ${cpu_list} ${range}
        echo "Bind interrupts of ${nic} to cpu $cpu_list. Done."
        exit 0
    fi

    # 按NUMA列表绑核
    if [ -n "${node_list}" ]; then
        local converted_cpu_list # 影响$?结果

        converted_cpu_list=$(convert_node_list_to_cpu_list ${node_list})
        local ret=$?
        if [ $ret -ne 0 ]; then
            exit 1
        fi

        set_affinity ${nic} ${converted_cpu_list} ${range}
        echo "Bind interrupts of ${nic} to node $node_list. Done."
        exit 0
    fi
}

main "$@"
